#  将mediapipe结果的pose，righthand，lefthand
# 按照x1,y1,z1,x2,y2,z2……存储
# 一张图片为一行

import cv2
import mediapipe as mp
import os
import pandas as pd
import numpy as np
from feature_function import del_file, get_dist_feature

# -----------参数设置
delete_csv = True

# -------------图片路径
root='D:/work_dxx/'
# root = 'D:/work_dxx/video_to_image/'
# filename = 'all_images/'
# filename='video_images_6/'
filename='video_images_test_5/'
resultname = 'result_' + filename[0:-1]
print(f'resultname={resultname}')
imageresult = root + resultname + '/'
path_data = imageresult

#------------标签------------

label_i = filename[-2]  # 视频转图像标签
# print('picture={}'.format(imagename))


#---------删除已有csv、图片----------
if delete_csv:
    try:
        os.remove(resultname + '_xyz.csv')
        print(f' successfully delete {resultname}_xyz.csv')
    except:
        pass
    try:
        os.remove(resultname + '_xyz_feature.csv')
        print(f' successfully delete {resultname}_xyz_feature.csv')
    except:
        pass
    try:
        del_file(path_data)
        print(f'successfully delete data in {path_data}')
    except:
        pass

filepath = os.path.join(root, filename)
# print(filepath)
image_paths = list(map(lambda x: os.path.join(filepath, x), os.listdir(filepath)))
# begin_delete_csv=True

# --------------读取数据并写入
mp_drawing = mp.solutions.drawing_utils
mp_holistic = mp.solutions.holistic

for i in range(len(image_paths)):
    # for i in range(1):
    image_path = image_paths[i]
    # print(image_path)
    imagename = image_path.split('/')[-1][:-len(".jpg")]
    print('{}th imagename={}'.format(i, imagename))

    # label_i = imagename[-1]
    # print(f'label_i={label_i}')


    with mp_holistic.Holistic(min_detection_confidence=0.5, min_tracking_confidence=0.5) as holistic:
        # holistic = mp_holistic.Holistic(static_image_mode=True)
        image = cv2.imread(image_path)
        image_hight, image_width, _ = image.shape
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        results = holistic.process(image)
        # 在图片上标记点
        mp_drawing.draw_landmarks(image, results.face_landmarks, mp_holistic.FACEMESH_CONTOURS,
                                  mp_drawing.DrawingSpec(color=(80, 110, 10), thickness=1, circle_radius=1),
                                  mp_drawing.DrawingSpec(color=(80, 256, 121), thickness=1, circle_radius=1))
        mp_drawing.draw_landmarks(image, results.pose_landmarks, mp_holistic.POSE_CONNECTIONS,
                                  mp_drawing.DrawingSpec(color=(80, 22, 10), thickness=2, circle_radius=4),
                                  mp_drawing.DrawingSpec(color=(80, 44, 121), thickness=2, circle_radius=2))
        # Right hand
        mp_drawing.draw_landmarks(image, results.right_hand_landmarks, mp_holistic.HAND_CONNECTIONS,
                                  mp_drawing.DrawingSpec(color=(121, 22, 76), thickness=2, circle_radius=4),
                                  mp_drawing.DrawingSpec(color=(121, 44, 250), thickness=2, circle_radius=2))

        # Left hand
        mp_drawing.draw_landmarks(image, results.left_hand_landmarks, mp_holistic.HAND_CONNECTIONS,
                                  mp_drawing.DrawingSpec(color=(245, 117, 66), thickness=2, circle_radius=4),
                                  mp_drawing.DrawingSpec(color=(245, 66, 230), thickness=2, circle_radius=2))



        if label_i in ['0','1','2'] and (
                (results.right_hand_landmarks==None
                and results.left_hand_landmarks==None)
            or (results.right_hand_landmarks != None
                and results.left_hand_landmarks != None)
            or (results.pose_landmarks.landmark[15].y < 1
                 and results.pose_landmarks.landmark[16].y < 1)
        ):
            os.remove(image_path)
            print(f'成功删除{label_i}图像{imagename}')
            continue
        elif label_i in ['4','5'] \
            and (results.right_hand_landmarks==None \
                or results.left_hand_landmarks==None):
            os.remove(image_path)
            print(f'成功删除{label_i}图像{imagename}')
            continue
        elif label_i in ['6'] \
            and (results.right_hand_landmarks!=None \
                or results.left_hand_landmarks!=None):
            os.remove(image_path)
            print(f'成功删除{label_i}图像{imagename}')
            continue
        elif label_i in ['3'] \
            and (results.pose_landmarks.landmark[15].y>1
                and results.pose_landmarks.landmark[16].y>1):
            os.remove(image_path)
            print(f'成功删除{label_i}图像{imagename}')
            continue
        # 读取坐标值
        if not os.path.exists(imageresult):
            os.makedirs(imageresult)
        cv2.imwrite(imageresult + '{}_result.jpg'.format(imagename), image)
        label = [int(label_i)]  # 标签
        poseData, righthandData, lefthandData, = [], [], []

        # poseData
        try:
            for i in range(len(results.pose_landmarks.landmark)):
                landmark = results.pose_landmarks.landmark[i]
                poseData += [landmark.x, landmark.y, landmark.z]
        except:
            poseData += [0] * 33 * 3
        #     print('pose_landmarks not exist!!! zeros{}'.format(33 * 2))
        # print(' after pose datalist length={}'.format(len(poseData)))
        # righthandData
        try:
            for i in range(len(results.right_hand_landmarks.landmark)):
                landmark = results.right_hand_landmarks.landmark[i]
                righthandData += [landmark.x, landmark.y, landmark.z]
        except:
            righthandData += [0] * 21 * 3
        #     print('right_hand_landmarks not exist!!! zeros{}'.format(21 * 2))
        # print(' after righthand datalist length={}'.format(len(righthandData)))
        # lefthandData
        try:
            for i in range(len(results.left_hand_landmarks.landmark)):
                landmark = results.left_hand_landmarks.landmark[i]
                lefthandData += [landmark.x, landmark.y, landmark.z]
        except:
            lefthandData += [0] * 21 * 3
        # dist_feature = label + get_dist_feature(poseData, righthandData, lefthandData)
        # print(f'dist_feature={dist_feature}')
        dataList = label + poseData + righthandData + lefthandData
        # print(f'len(dataList)={len(dataList)}')
        # print(f'dist_feature={dist_feature}')

    # 若表不存在则创建，若存在则写入数据
        if not os.path.exists(resultname + '_xyz.csv'):
            poseName = ['label']
            for i in range(33):
                poseName += ['x_pose' + str(i), 'y_pose' + str(i), 'z_pose' + str(i)]
            rightHandName = []
            for i in range(21):
                rightHandName += ['x_righthand' + str(i), 'y_righthand' + str(i), 'z_righthand' + str(i)]
            leftHandName = []
            for i in range(21):
                leftHandName += ['x_lefthand' + str(i), 'y_lefthand' + str(i), 'z_lefthand' + str(i)]
            name = poseName + rightHandName + leftHandName
            df = pd.DataFrame(columns=name)
            df.to_csv(resultname + '_xyz.csv', mode='w', index=False, sep=',')
            df = pd.DataFrame(dataList).T
            df.to_csv(resultname + '_xyz.csv', mode='a', index=False, header=False, sep=',')
        else:
            df = pd.DataFrame(dataList).T
            df.to_csv(resultname + '_xyz.csv', mode='a', index=False, header=False, sep=',')

        # if not os.path.exists(resultname + '_xyz_feature.csv'):
        #     featurename = ['label']
        #     for i in range(len(dist_feature) - 1):
        #         featurename += ['feature' + str(i)]
        #     df = pd.DataFrame(columns=featurename)
        #     df.to_csv(resultname + '_xyz_feature.csv', mode='w', index=False, sep=',')
        #     df = pd.DataFrame(dist_feature).T
        #     df.to_csv(resultname + '_xyz_feature.csv', mode='a', index=False, header=False, sep=',')
        # else:
        #     df = pd.DataFrame(dist_feature).T
        #     df.to_csv(resultname + '_xyz_feature.csv', mode='a', index=False, header=False, sep=',')
